# src/database_connector.py
import json
from typing import Dict, Any, Optional, List

import mysql.connector


class DatabaseConnector:
    """
    连接到源数据库并提取其 Schema 信息，包括表和字段的注释。
    """

    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.db_name = config['database']
        self.conn = mysql.connector.connect(
            host=config['host'],
            user=config['user'],
            port=config['port'],
            password=config['password'],
            database='information_schema'
        )

    def get_schema(self) -> Dict[str, Dict[str, Any]]:
        """
        获取指定数据库的所有表（包括注释）和列（包括注释）的元数据。
        返回结构: {'table_name': {'table_comment': '...', 'columns': [...]}}
        """
        schema = {}
        cursor = self.conn.cursor(dictionary=True)

        # 1. 查询所有表及其注释
        cursor.execute("""
            SELECT LOWER(table_name) AS table_name, LOWER(table_comment) AS table_comment
            FROM tables
            WHERE table_schema = %s
        """, (self.db_name,))
        tables = cursor.fetchall()

        # 2. 为每个表查询其所有列及其注释
        for table in tables:
            table_name = table['table_name']
            table_comment = table['table_comment']

            cursor.execute("""
                SELECT
                    LOWER(column_name) AS column_name,
                    LOWER(data_type) AS data_type,
                    LOWER(column_comment) AS column_comment,
                    (COLUMN_KEY = 'PRI') AS is_primary_key
                FROM columns
                WHERE table_schema = %s AND table_name = %s
            """, (self.db_name, table_name))
            columns = cursor.fetchall()

            schema[table_name] = {
                "table_comment": table_comment,
                "columns": columns
            }

        cursor.close()
        return schema

    def execute_query(self, sql_query: str) -> tuple[list, list]:
        """
        在源数据库上执行一个查询并返回表头和所有数据行。
        """

        conn = mysql.connector.connect(
            host=self.config['host'],
            user=self.config['user'],
            port=self.config['port'],
            password=self.config['password'],
            database=self.config['database']
        )
        cursor = conn.cursor(dictionary=True)
        cursor.execute(sql_query)

        columns = [i[0] for i in cursor.description]
        data = cursor.fetchall()

        conn.close()
        return columns, data


class MetadataRepository:
    """
    将丰富的元数据写入到系统的元数据库中。
    """

    def __init__(self, config: Dict[str, Any]):
        self.conn = mysql.connector.connect(**config)

    def save_table(self, table_data: Dict[str, Any]) -> int:
        """
        保存表元数据，如果已存在则更新。返回表的ID。
        """
        cursor = self.conn.cursor()
        sql = """
            INSERT INTO tables (table_name, semantic_name, description)
            VALUES (%(table_name)s, %(semantic_name)s, %(description)s)
            ON DUPLICATE KEY UPDATE
                semantic_name = VALUES(semantic_name),
                description = VALUES(description)
        """
        cursor.execute(sql, table_data)
        self.conn.commit()

        # 获取刚插入或更新的记录的ID
        cursor.execute("SELECT id FROM tables WHERE table_name = %s", (table_data['table_name'],))
        result = cursor.fetchone()
        cursor.close()
        return result[0] if result else None

    def save_column(self, column_data: Dict[str, Any]) -> int:
        """
        保存列元数据，如果已存在则更新。返回列的ID。
        """
        cursor = self.conn.cursor()
        sql = """
            INSERT INTO columns (table_id, column_name, data_type, semantic_name, description, is_primary_key)
            VALUES (%(table_id)s, %(column_name)s, %(data_type)s, %(semantic_name)s, %(description)s, %(is_primary_key)s)
            ON DUPLICATE KEY UPDATE
                data_type = VALUES(data_type),
                semantic_name = VALUES(semantic_name),
                description = VALUES(description),
                is_primary_key = VALUES(is_primary_key)
        """
        cursor.execute(sql, column_data)
        self.conn.commit()

        cursor.execute("SELECT id FROM columns WHERE table_id = %s AND column_name = %s",
                       (column_data['table_id'], column_data['column_name']))
        result = cursor.fetchone()
        cursor.close()
        return result[0] if result else None

    def get_schema_as_string(self) -> str:
        """
        从元数据库中查询所有表和列的信息，并格式化为对LLM友好的字符串。
        """
        cursor = self.conn.cursor(dictionary=True)

        # 查询所有表
        cursor.execute("SELECT id, table_name, semantic_name, description FROM tables")
        tables = cursor.fetchall()

        schema_parts = []
        for table in tables:
            table_str = f"Table `{table['table_name']}` (semantic_name: {table['semantic_name']}, description: {table['description']}):\n"

            # 查询该表的所有列
            cursor.execute(
                "SELECT column_name, data_type, semantic_name, description FROM columns WHERE table_id = %s",
                (table['id'],)
            )
            columns = cursor.fetchall()

            for col in columns:
                table_str += f"  - `{col['column_name']}` ({col['data_type']}) (semantic_name: {col['semantic_name']}, description: {col['description']})\n"

            schema_parts.append(table_str)

        cursor.close()
        return "\n".join(schema_parts)

    def close(self):
        if self.conn.is_connected():
            self.conn.close()


class TemplateRepository:
    """
    负责报表模板和参数的数据库操作。
    """

    def __init__(self, config: Dict[str, Any]):
        self.conn = mysql.connector.connect(**config)

    def save_template(self, name: str, description: str, example_descriptions: Optional[List[str]],
                      template_file_path: str) -> int:
        """
        保存报表模板主信息，包括JSON格式的用例描述。
        """
        cursor = self.conn.cursor()
        # 将列表转换为JSON字符串进行存储
        example_desc_json = json.dumps(example_descriptions) if example_descriptions else None

        sql = """
            INSERT INTO report_templates (template_name, description, example_descriptions, example_file_path, status)
            VALUES (%s, %s, %s, %s, 'Active')
        """
        cursor.execute(sql, (name, description, example_desc_json, template_file_path))
        self.conn.commit()
        template_id = cursor.lastrowid
        cursor.close()
        return template_id

    def get_template_by_id(self, template_id: int) -> dict | None:
        """
        根据ID从数据库中获取单个主模板。
        """
        cursor = self.conn.cursor(dictionary=True)
        sql = "SELECT * FROM report_templates WHERE id = %s"
        cursor.execute(sql, (template_id,))
        result = cursor.fetchone()
        cursor.close()
        return result

    def save_mapped_parameters(self, template_id: int, mappings: List[Dict[str, Any]]):
        """
        批量保存模板字段到数据库列的映射关系。
        """
        if not mappings:
            return

        cursor = self.conn.cursor()
        sql = """
            INSERT INTO report_parameter (template_id, template_field, db_column, description)
            VALUES (%s, %s, %s, %s)
        """
        data_to_insert = [
            (template_id, m['template_field'], m['db_column'], m.get('description', ''))
            for m in mappings
        ]
        cursor.executemany(sql, data_to_insert)
        self.conn.commit()
        cursor.close()

    def get_mapped_parameters(self, template_id: int) -> List[Dict[str, Any]]:
        """
        获取一个模板所有已映射的参数。
        """
        cursor = self.conn.cursor(dictionary=True)
        sql = "SELECT template_field, db_column, description FROM report_parameter WHERE template_id = %s"
        cursor.execute(sql, (template_id,))
        results = cursor.fetchall()
        cursor.close()
        return results

    def get_all_active_templates(self) -> List[Dict]:
        """
        获取所有状态为 'Active' 的模板，用于向量化。
        """
        cursor = self.conn.cursor(dictionary=True)
        sql = "SELECT id, template_name, description, example_descriptions FROM report_templates WHERE status = 'Active'"
        cursor.execute(sql)
        results = cursor.fetchall()
        cursor.close()
        return results


class QueryCacheRepository:
    """
    负责语义查询缓存的数据库操作。
    """

    def __init__(self, config: Dict[str, Any]):
        self.conn = mysql.connector.connect(**config)

    def save_query_pair(self, nlq: str, sql: str, template_id: int) -> int:
        """
        保存一个“自然语言-SQL-模板ID”三元组。
        """
        cursor = self.conn.cursor()
        query = """
            INSERT INTO query_cache (natural_language_query, generated_sql, template_id)
            VALUES (%s, %s, %s)
        """
        cursor.execute(query, (nlq, sql, template_id))
        self.conn.commit()
        cache_id = cursor.lastrowid
        cursor.close()
        return cache_id

    def get_by_id(self, cache_id: int) -> dict | None:
        """
        根据ID获取一条缓存记录。
        """
        cursor = self.conn.cursor(dictionary=True)
        cursor.execute("SELECT * FROM query_cache WHERE id = %s", (cache_id,))
        result = cursor.fetchone()
        cursor.close()
        return result

    def close(self):
        if self.conn.is_connected():
            self.conn.close()
